import torch
import torch.nn as nn
import math  
import torch.nn.functional as F

class PositionalEncoding(nn.Module):
    # 位置编码类
    def __init__(self, dimModel, dropout=0.1, maxLen=1024):
        """
        初始化位置编码模块。
        :param dimModel: 模型维度。
        :param dropout: Dropout 比率。
        :param maxLen: 最大序列长度。
        """
        super(PositionalEncoding, self).__init__()
        self.dropout = nn.Dropout(p=dropout)

        # 创建位置编码
        pe = torch.zeros(maxLen, dimModel)
        position = torch.arange(0, maxLen, dtype=torch.float).unsqueeze(1)
        divTerm = torch.exp(torch.arange(0, dimModel, 2).float() * (-math.log(10000.0) / dimModel))
        pe[:, 0::2] = torch.sin(position * divTerm)
        pe[:, 1::2] = torch.cos(position * divTerm)
        pe = pe.unsqueeze(0).transpose(0, 1)
        self.register_buffer('pe', pe)

    def forward(self, x):
        """
        应用位置编码。
        :param x: 输入张量，形状为 [seq_len, batch_size, dModel]。
        :return: 应用了位置编码后的张量。
        """
        x = x + self.pe[:x.size(0), :]
        return self.dropout(x)
    
class TransformerModel(nn.Module):
    # Transformer 模型
    def __init__(self, modelDim, numHeads, numEncoderLayers, dimFeedforward, maxLen=1024, dropout=0.1):
        """
        初始化 Transformer 模型。
        :param modelDim: 模型维度。
        :param numHeads: 注意力头的数量。
        :param numEncoderLayers: 编码器层数。
        :param dimFeedforward: 前馈网络维度。
        :param maxLen: 最大序列长度。
        :param dropout: Dropout 比率。
        """
        
        super(TransformerModel, self).__init__()
        self.positionalEncoding = PositionalEncoding(modelDim, dropout=dropout, maxLen=maxLen)
        self.modelDim = modelDim
        self.maxLen = maxLen
        
        encoderLayers = nn.TransformerEncoderLayer(
            d_model=modelDim, 
            nhead=numHeads, 
            dim_feedforward=dimFeedforward, 
            dropout=dropout,
            batch_first = True
        )
        self.transformerEncoder = nn.TransformerEncoder(encoderLayers, num_layers=numEncoderLayers)
        self.gelu = nn.GELU()
        self.dropout = nn.Dropout(dropout)

    def forward(self, src, paddingMask):
        """
        前向传播。
        :param src: 源序列张量，形状为 [batch_size,seq_len,modelDim]。
        :param paddingMask: 填充掩码张量，形状为 [seq_len, batch_size](transformerEncoder要求这样)，指示填充位置。
        :return: Transformer 编码器的输出。
        """
        paddingMask = paddingMask.transpose(0,1)
        src = self.positionalEncoding(src)
        #Transformer对于paddingmask的定义是，True为填充位置，需要忽略，而我们数据处理中的paddingmask却是false为填充位置，因此这里取反
        src = self.transformerEncoder(src, src_key_padding_mask=~paddingMask)
        src = self.gelu(src)
        output = self.dropout(src)
        # #这里transpose是因为paddingMask的形状原本为[seq_len, batch_size]，要将它转回[batch_size,seq_len]
        # output = output * paddingMask.transpose(0, 1).unsqueeze(-1).expand(-1, -1, self.modelDim)
        return output

class LinearMapping(nn.Module):
    # 线性映射层
    def __init__(self, inputDim, outputDim, dropout=0.0):
        """
        初始化线性映射层。
        :param inputDim: 输入维度。
        :param outputDim: 输出维度。
        :param dropout: Dropout 比率。
        """
        super(LinearMapping, self).__init__()
        self.inputDim = inputDim
        self.outputDim = outputDim
        self.inputLinear = nn.Linear(inputDim, outputDim)
        self.dropout = nn.Dropout(p=dropout)

    def forward(self, src):
        """
        前向传播。
        :param src: 源序列张量，形状应为 [batchsize, length, inputDim]。
        :return: 线性映射后的输出张量。
        """
        batchsize, length, _ = src.size()
        src = src.reshape(batchsize * length, -1)
        src = self.inputLinear(src)
        output = self.dropout(src)
        output = output.view(batchsize, length, -1)
        return output

class PreTrainModel(nn.Module):
    # 预训练模型
    def __init__(self, classes_num,inputLayer:LinearMapping, mainTransformerModel: TransformerModel):
        """
        初始化预训练模型。
        :param inputLayer: 输入线性映射层。这个部分需要根据输入数据的形状而确定，因此对于不同数据集这个部分需要不断改变，且没有激活层
        :param mainTransformerModel: 主 Transformer 模型。这个部分是作为通用模型连续训练，因此这个部分的参数需要保存
        :param outputLayer: 输出线性映射层。这个部分需要根据输出数据的形状而确定，因此对于不同数据集这个部分需要不断改变，且没有激活层
        """
        super(PreTrainModel, self).__init__()
        self.inputLayer = inputLayer
        self.mainTransformerModel = mainTransformerModel
        self.bn1 = nn.BatchNorm1d(mainTransformerModel.maxLen)
        # self.bn2 = nn.BatchNorm1d(outputLayer.outputDim)
        self.lastLayer = nn.Linear(mainTransformerModel.maxLen*mainTransformerModel.modelDim, classes_num)
        self.activate1 = nn.GELU()   
        self.activate2 = nn.GELU()   
        self.softmax = nn.Softmax(dim=-1)
        
    def forward(self, src, paddingMask):
        """
        前向传播。
        :param src: 源序列张量，形状应为 [batchsize, length, inputDim]。
        :param paddingMask: 填充掩码张量，形状为 [seq_len, batch_size]，指示填充位置。
        :return: 模型的输出。
        """
        src = self.inputLayer(src) * math.sqrt(self.inputLayer.outputDim)
        src = self.mainTransformerModel(src, paddingMask)#[batchsize,Maxlen,d_model]
        #将nan值转换为0
        src = torch.nan_to_num(src)

        src = self.bn1(src)
        src = self.activate1(src)#[batchsize,Maxlen,d_model]

        src = src.reshape(-1,self.mainTransformerModel.maxLen*self.mainTransformerModel.modelDim)#[batchsize,-1]
        # src = src.squeeze(2)    # [batch_size, sequence_length]
        src = self.lastLayer(src)
        # print("全连接层输出形状:",src.shape)  
    
        output = self.softmax(src)
        return output
        

# # 创建模型实例和随机数据，跑通整个前向传播过程
# batchsize = 10
# length = 30
# inputDim = 20
# outputDim = 20
# modelDim = 32
# numHeads = 16
# numEncoderLayers = 4
# dimFeedforward = 45
# dropout = 0.1

# inputData = torch.randn(batchsize, length, inputDim)
# paddingMask = torch.ones((batchsize, length), dtype=torch.bool)
# paddingMask[:, -5:] = False
# paddingMask = paddingMask.transpose(0, 1)

# inputLinear = LinearMapping(inputDim, modelDim)
# mainModel = TransformerModel(modelDim=32, numHeads=16, numEncoderLayers=6, dimFeedforward=128, dropout=0.3, maxLen=4096)
# outputLinear = LinearMapping(modelDim, outputDim)
# model = PreTrainModel(inputLayer=inputLinear, mainTransformerModel=mainModel, outputLayer=outputLinear)
# outputData = model(inputData, paddingMask)
# print(outputData.shape)
# total = sum([param.nelement() for param in model.parameters()])
# print(total)
